import struct

# Constants
IMA_ADPCM_DVI4 = 0
IMA_ADPCM_VDVI = 1

# IMA ADPCM state descriptor
class IMAADPCMState:
    def __init__(self, variant):
        self.variant = variant
        self.last = 0
        self.step_index = 0
        self.ima_byte = 0
        self.bits = 0

# Helper function to saturate values to int16 range
def saturate(amp):
    return max(min(amp, 32767), -32768)

# Decode ADPCM code to linear PCM
def decode(s, adpcm):
    step_size = [7, 8, 9, 10, 11, 12, 13, 14, 16, 17, 19, 21, 23, 25, 28, 31,
                 34, 37, 41, 45, 50, 55, 60, 66, 73, 80, 88, 97, 107, 118, 130,
                 143, 157, 173, 190, 209, 230, 253, 279, 307, 337, 371, 408, 449,
                 494, 544, 598, 658, 724, 796, 876, 963, 1060, 1166, 1282, 1411,
                 1552, 1707, 1878, 2066, 2272, 2499, 2749, 3024, 3327, 3660, 4026,
                 4428, 4871, 5358, 5894, 6484, 7132, 7845, 8630, 9493, 10442, 11487,
                 12635, 13899, 15289, 16818, 18500, 20350, 22385, 24623, 27086,
                 29794, 32767]

    e = step_size[s.step_index] >> 3
    if adpcm & 0x01:
        e += step_size[s.step_index] >> 2
    if adpcm & 0x02:
        e += step_size[s.step_index] >> 1
    if adpcm & 0x04:
        e += step_size[s.step_index]
    if adpcm & 0x08:
        e = -e

    linear = saturate(s.last + e)
    s.last = linear
    s.step_index += [0, -1, -1, -1, 2, 4, 6, 8][adpcm & 0x07]
    s.step_index = max(min(s.step_index, 88), 0)
    return linear

# Encode linear PCM to ADPCM code
def encode(s, linear):
    step_size = [7, 8, 9, 10, 11, 12, 13, 14, 16, 17, 19, 21, 23, 25, 28, 31,
                 34, 37, 41, 45, 50, 55, 60, 66, 73, 80, 88, 97, 107, 118, 130,
                 143, 157, 173, 190, 209, 230, 253, 279, 307, 337, 371, 408, 449,
                 494, 544, 598, 658, 724, 796, 876, 963, 1060, 1166, 1282, 1411,
                 1552, 1707, 1878, 2066, 2272, 2499, 2749, 3024, 3327, 3660, 4026,
                 4428, 4871, 5358, 5894, 6484, 7132, 7845, 8630, 9493, 10442, 11487,
                 12635, 13899, 15289, 16818, 18500, 20350, 22385, 24623, 27086,
                 29794, 32767]

    e = linear - s.last
    if e < 0:
        adpcm = 0x08
        e = -e
    else:
        adpcm = 0x00

    delta = step_size[s.step_index] >> 3
    code = 0
    for i in range(3):
        if e >= delta:
            code |= 0x04 >> i
            e -= delta
        delta >>= 1

    if adpcm & 0x08:
        e = -e

    s.last = saturate(s.last + e)
    s.step_index += [0, -1, -1, -1, 2, 4, 6, 8][adpcm & 0x07]
    s.step_index = max(min(s.step_index, 88), 0)
    return adpcm | code

# Encode linear PCM data from input file to ADPCM and write to output file
def encode_file(input_file, output_file, variant=IMA_ADPCM_DVI4):
    # Initialize ADPCM state
    state = IMAADPCMState(variant)

    # Open input file for reading binary data
    with open(input_file, 'rb') as input_f:
        # Open output file for writing binary data (creating if not exists, overwrite if exists)
        with open(output_file, 'wb') as output_f:
            while True:
                # Read a chunk of PCM data (arbitrary size, adjust as needed)
                pcm_data = input_f.read(1024)
                if not pcm_data:
                    # Reached end of file
                    break

                # Convert binary data to list of 16-bit PCM samples
                pcm_values = struct.unpack('h' * (len(pcm_data) // 2), pcm_data)

                # Encode PCM data to ADPCM
                adpcm_data = bytearray()
                for sample in pcm_values:
                    adpcm_sample = encode(state, sample)
                    adpcm_data.append(adpcm_sample)

                # Write encoded ADPCM data to output file
                output_f.write(adpcm_data)

# Decode ADPCM data from input file to linear PCM and write to output file
def decode_file(input_file, output_file, variant=IMA_ADPCM_DVI4):
    # Initialize ADPCM state
    state = IMAADPCMState(variant)

    # Open input file for reading binary data
    with open(input_file, 'rb') as input_f:
        # Open output file for writing binary data (creating if not exists, overwrite if exists)
        with open(output_file, 'wb') as output_f:
            while True:
                # Read a chunk of ADPCM data (arbitrary size, adjust as needed)
                adpcm_data = input_f.read(1024)
                if not adpcm_data:
                    # Reached end of file
                    break

                # Decode ADPCM data to PCM
                pcm_data = bytearray()
                for adpcm_sample in adpcm_data:
                    pcm_sample = decode(state, adpcm_sample)
                    pcm_data.extend(struct.pack('h', pcm_sample))

                # Write decoded PCM data to output file
                output_f.write(pcm_data)


# Example usage
if __name__ == "__main__":
    encode_file("input.pcm", "output.adpcm")
    decode_file("output.adpcm", "decoded_output.pcm")
